# aem.two_d.transient.constant_aquifer
#
# Copyright (c) 2000, V.Kelson and M.Bakker
#           (c) 2001, Willem Jan Zaadnoordijk

import copy
import math

import geometry
SearchableList = geometry.two_d.collections.SearchableList
XY = geometry.XY

import aem.container
ColumnList=aem.container.ColumnList
RowList=aem.container.RowList
TraceLine=aem.container.TraceLine

import aem.matrix

import base_aquifer

class ConstantAquifer(base_aquifer.Aquifer):
	"""Aquifer class with constant properties, a reference point and uniform flow

	Python call: ConstantAquifer(b = 0.0,H = 1.0,k = 1.0,n = 0.2, 
		     S = 0.1, Haverage = -99.0,
		     refPt = geometry.two_d.Point(0,0),refHead = 1.0, 
		     uniformFlow = geometry.XY(0,0),parent=None)

	Attributes:

	  - b: base elevation of aquifer
	  - H: thickness of aquifer
	  - k: hydraulic conductivity of aquifer
	  - n: porosity of aquifer
	  - S: storativity (defaults to 0.1)
	  - Haverage: average saturated thickness (defaults to H)
	  - refPt: location of reference point, XY object
	  - refHead: head at reference point
	  - uniformFlow: Qx,Qy components of uniform flow, XY object
	  - topElevation: elevation of top of aquifer
	  - topPotential: potential corresponding to head at top of aquifer
	  - parent: parent of the element, set to None (inherited from Element)
	  - parameters: ColumnList of strength parameters, to list of three entries,
	    the constant, and the components of uniform flow Qx and Qy
	  - solver: solver (inherited from Element)
	  - elementList: list of elements, contains all elements in aquifer (inherited from Element)

	Overloaded functions from Element: __init__, potentialInfluence, dischargeInfluence,
	getMatrixCoefficients, getMatrixContribution, check

	Overloaded from Aquifer: headToPotential, potentialToHead, transmissivity
	"""
	def __init__(self,b = 0.0,H = 1.0,k = 1.0,n = 0.2, \
		     S = 0.1, Haverage = -99.0,            \
		     refPt = geometry.two_d.Point(0,0),refHead = 1.0, \
		     uniformFlow = geometry.XY(0,0),parent=None):
		'''Constructor
		Accepts the aquifer properties and reference, reference point and uniform flow
		'''
		base_aquifer.Aquifer.__init__(self)
		self.S=S
		self.Haver=Haverage
		if Haverage < 0 :
			self.kH_S = k * H / S
		else:
			self.kH_S = k * Haverage / S
		self.initialSteadyState=aem.two_d.single.ConstantAquifer( \
			b,H,k,n,refPt,refHead,uniformFlow,None)
		self.parameters = ColumnList([])

	def __repr__(self):
		return 'ConstantAquifer{kH/S=' + str((self.kH_S))+ \
		   '; steadyAq=' + str((self.initialSteadyState))  \
		   + '}'

	def potentialInfluence(self,xy,t=0):
		return RowList( [] )
	def dischargeInfluence(self,xy,t=0):
		return RowList( [] )
	def rechargeInfluence(self,xy,t=0):
		return RowList( [] )
	def extractionInfluence(self,t=0):
		return RowList( [] )

	def getMatrixCoefficients(self):
		return [ ]
	def getMatrixContribution(self,elementList):
		# Returns the MatrixRow objects for this element
		return ( [],[] )
	def check(self):
		return [ ]

	def headToPotential(self,pt,head):
		#Convert head to potential at the point pt
		return self.initialSteadyState.headToPotential(pt,head)

	def potentialToHead(self,pt,pot):
		#Convert head to potential at the point pt
		return self.initialSteadyState.potentialToHead(pt,head)

	def transmissivity(self,pt,t=0):
		#transmissivity at the point pt
		return self.initialSteadyState.transmissivity(pt,t)

	def velocity(self,xyzpt,t=0):
		pt = xyzpt.XY()
		QxQy = self.discharge(pt,t)
		h = self.head(pt,t)
		if h < self.initialSteadyState.topElevation:
			qxqy = QxQy / (h - self.initialSteadyState.b)
			qz = -( self.recharge(pt,t) + \
				( QxQy.lengthSquared() )/( self.initialSteadyState.k * h **2 ) ) * \
				( xyzpt.z - self.initialSteadyState.b ) / h
		else:
			qxqy = QxQy / self.initialSteadyState.H
			qz = - self.recharge(pt,t) * \
			   (xyzpt.z - self.initialSteadyState.b) / self.initialSteadyState.H
		rv = geometry.XYZ(qxqy.x,qxqy.y,qz) / self.initialSteadyState.n
		return rv

	def trace(self,xyzstart,step,tmax,maxsteps=10,t=0):
		'''Tracing routine

		Input:
		- xyzstart: XYZ object
		- step: 3D step size
		- tmax: maximum time (will stop at tmax if no other element has been reached)
		- maxsteps: maximum number of steps (default = 10)
		- t: starting time

		Returns:
		- TraceLine object'''

		# Check starting point
		h = self.head(xyzstart.XY(),t)
		if xyzstart.z < self.initialSteadyState.b or \
		   xyzstart.z > min(h,self.initialSteadyState.topElevation):
			raise 'Starting point outside saturated aquifer'

		traceLine = TraceLine()
		traceLine = traceLine.addpoint(xyzstart,t)

		xyzold = xyzstart
		told = t
		nstep= 0

		while nstep < maxsteps:
			# Compute velocity and stepsize
			vxyz = self.velocity(xyzold,told)
			velo = vxyz.magnitude()
			if velo==0: break
			tstep = step / velo
			# Compute new point
			xyznew = xyzold + vxyz * tstep
			tnew = told + tstep
			changed = 0
			# Check whether point is at or near element
			for e in self.initialSteadyState.elementList:
				check = e.nearElement(xyzold,xyznew,step)
				if check.change:
					xyznew = check.xyzchanged
					stepnew = (xyznew-xyzold).magnitude()
					tnew = told + stepnew/step * tstep + check.delt
					changed = 1
				if check.stop: break
			for e in self.elementList:
				check = e.nearElement(xyzold,xyznew,step)
				if check.change:
					xyznew = check.xyzchanged
					stepnew = (xyznew-xyzold).magnitude()
					tnew = told + stepnew/step * tstep + check.delt
					changed = 1
				if check.stop: break
			# Check whether point is inside aquifer
			# Does step out somewhat crudely, since it may step above water table
			if self.recharge(xyzold.XY(),told) < 0:
				if xyznew.z > min( self.initialSteadyState.topElevation, \
				                   self.head(xyznew.XY(),tnew) ):
					check.stop = 1
			# If no change, do corrector step
			if not changed and not check.stop:
				vxyznew = self.velocity(xyznew,tnew)
				vxyz = 0.5 * (vxyz+vxyznew)
				velo = vxyz.magnitude()
				tstep = step / velo
				xyznew = xyzold + vxyz * tstep
				tnew = told + tstep
				# Check again whether is at or near point
				for e in self.initialSteadyState.elementList:
					check = e.nearElement(xyzold,xyznew,step)
					if check.change:
						xyznew = check.xyzchanged
						stepnew = (xyznew-xyzold).magnitude()
						tnew = told + stepnew/step * tstep + check.delt
					if check.stop: break
				for e in self.elementList:
					check = e.nearElement(xyzold,xyznew,step)
					if check.change:
						xyznew = check.xyzchanged
						stepnew = (xyznew-xyzold).magnitude()
						tnew = told + stepnew/step * tstep + check.delt
					if check.stop: break
			# Check whether point is inside aquifer
			# Does step out somewhat crudely, since it may step above water table
			if self.recharge(xyzold.XY(),told) < 0:
				if xyznew.z > min( self.initialSteadyState.topElevation, self.head(xyznew.XY()) ):
					check.stop = 1
			# Check whether maximum time has been reached
			if tnew > tmax:
				xyznew = xyzold + (tmax-told) * vxyz
				tnew = tmax
				check.stop = 1
			# Add point to trace line
			traceLine = traceLine.addpoint(xyznew,tnew)
			nstep = nstep + 1
			if check.stop: break
			xyzold=xyznew; told=tnew

		if velo == 0:
			print "Stagnation point"
		else:
			if tnew == tmax:
				print "Reached maximum time"
			else:
				if nstep >= maxsteps:
					print "Reached maximum number of steps"
				else:
					if check.stop and check.change:
						print "Reached element"

		return traceLine

	def tracelines(self,xyz,step,tmax,Nmax,filename,t=0):
		'''Routine for calculating multiple tracelines and writing them to file.
		xyz is list of XYZ objects, filename is character string, such as 'mark.dat'
		and must be between quotes
		'''
		# clear file
		out = open(filename,'w')
		out.write('\n')
		out.close()
		for i in range(len(xyz)):
			if len(xyz)==len(t):
				tstart=t[i]
			else:
				tstart=t
			trace=self.trace(xyz[i],step,tmax,Nmax,tstart)
			trace.appendtofile(filename)

	def makegrid(self,xmin,xmax,xstep,ymin,ymax,ystep,filename,t=0):
		'''Routine for calculating grid of heads and write surfer grid
		filename must be character string, such as /temp/tim.grd (forward slashes!)'''
		nx=int((xmax-xmin)/xstep)
		ny=int((ymax-ymin)/ystep)
		matrix=[]
		zmin = self.head(XY(xmin,ymin)); zmax = zmin
		for j in range(ny+1):
			row = []
			y=ymin + j*ystep
			for i in range(nx+1):
				x=xmin + i*xstep;
				h = self.head(XY(x,y),t)
				if h < zmin: zmin = h
				if h > zmax: zmax = h
				row = row + [ h ]
			matrix = matrix + [row]

		out = open(filename,'w')
		out.write('DSAA\n')

		out.write(str(nx+1)+' '+str(ny+1)+'\n')
		out.write(str(xmin)+' '+str(xmin+nx*xstep)+'\n')
		out.write(str(ymin)+' '+str(ymin+ny*ystep)+'\n')
		out.write(str(zmin)+' '+str(zmax)+'\n')

		for j in range(ny+1):
			row = matrix[j]
			for i in range(nx):
				out.write(str(row[i])+' ')
			out.write(str(row[nx])+'\n')
		out.close()

	def makeAdo(self,fileUng,fileAdo,t=0):
		'''Routine for generating Adore file with heads
		at the points provided in an (TESNOD.UNG) Ascii
		file with the ArcInfo Ungenerate format
		filenames must be character strings,
		such as /temp/tim.grd (forward slashes!)'''

		fin=open(fileUng,'r')
		AdoSet=[]
		iend=1
		while iend>-1:
			aline=fin.readline()
			iend=len(aline)-1
			if aline.find('END')!=-1:
				iend=-99
			if iend>0:
				i=0
				aline=aline.lstrip()
				an,ax,ay=aline.split(None,2)
				x=float(ax)
				y=float(ay)
				h = self.head(XY(x,y),t)
				AdoSet = AdoSet + [ h ]

		fin.close()

		iend = len(AdoSet)
		if (iend>0):
			fout = open(fileAdo,'w')
			aline='*SET*head,t='+str(t)+'\n'
			fout.write(aline)
			fout.write('2\n')
			fout.write(str(iend).rjust(5)+'     (1f20.0)\n')
			for i in range(iend):
				ah=str(AdoSet[i])
				aline=ah.rjust(20)+'\n'
				fout.write(aline)
			fout.write('ENDSET\n')
			fout.close()
